import os
import marshal
import struct
import hashlib
import types
from typing import Dict, Optional, List, Any
from dataclasses import dataclass

@dataclass
class PyFileInfo:
    """Python文件信息"""
    name: str
    offset: int
    size: int
    checksum: str

class PythonLoader:
    """Python代码加载器"""
    MAGIC = 0x44545059
    
    def __init__(self, bin_file: str):
        """初始化加载器"""
        self.bin_file = bin_file
        self.files: Dict[str, PyFileInfo] = {}
        self._load_index()
    
    def _calculate_checksum(self, data: bytes) -> str:
        """计算校验和"""
        return hashlib.sha256(data).hexdigest()
    
    def _load_index(self):
        """加载索引"""
        with open(self.bin_file, 'rb') as f:
            # 读取文件头
            magic = struct.unpack('<I', f.read(4))[0]
            if magic != self.MAGIC:
                raise ValueError("无效的Python资源文件")
            
            # 读取版本和文件数量
            version = struct.unpack('<I', f.read(4))[0]
            file_count = struct.unpack('<I', f.read(4))[0]
            
            # 读取文件校验和
            file_checksum = f.read(32).hex()
            
            # 验证文件完整性
            checksum_pos = f.tell()
            f.seek(12)  # 跳过魔数、版本号、文件数量
            data_for_checksum = f.read()
            if self._calculate_checksum(data_for_checksum) != file_checksum:
                raise ValueError("Python资源文件已被修改")
            
            # 读取文件索引
            f.seek(checksum_pos)
            for _ in range(file_count):
                # 读取文件名
                name_len = struct.unpack('<I', f.read(4))[0]
                name = f.read(name_len).decode('utf-8')
                
                # 读取文件信息
                offset = struct.unpack('<Q', f.read(8))[0]
                size = struct.unpack('<Q', f.read(8))[0]
                checksum = f.read(32).hex()
                
                # 保存文件信息
                self.files[name] = PyFileInfo(name, offset, size, checksum)
    
    def load_module(self, module_path: str) -> Optional[types.ModuleType]:
        """加载Python模块"""
        if module_path not in self.files:
            return None
        
        try:
            info = self.files[module_path]
            
            # 读取字节码
            with open(self.bin_file, 'rb') as f:
                f.seek(info.offset)
                bytecode = f.read(info.size)
            
            # 验证字节码完整性
            if self._calculate_checksum(bytecode) != info.checksum:
                raise ValueError(f"Python模块已被修改: {module_path}")
            
            # 反序列化字节码
            code = marshal.loads(bytecode)
            
            # 创建模块
            module_name = f"dynamic_{info.name.replace('/', '_').replace('.', '_')}"
            module = types.ModuleType(module_name)
            
            # 执行字节码
            exec(code, module.__dict__)
            return module
            
        except Exception as e:
            print(f"加载模块失败 {module_path}: {e}")
            return None
    
    def load_function(self, module_path: str, function_name: str) -> Optional[callable]:
        """加载Python函数"""
        module = self.load_module(module_path)
        if module and hasattr(module, function_name):
            return getattr(module, function_name)
        return None
    
    def list_modules(self) -> List[str]:
        """列出所有模块"""
        return list(self.files.keys())